﻿using System.Collections.Generic;
using UnityEngine;
using Mars.Pathfinding;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;

public class FlowFieldPathfinding : MonoBehaviour
{
    [Header("Grid Settings")]
    public int gridWidth = 50;
    public int gridHeight = 50;
    public float cellSize = 1f;
    public LayerMask obstacleLayer;

    [Header("Job Settings")]
    public bool useJobs = true;
    public int jobBatchSize = 64;
    public bool enableDynamicObstacles = true;
    public float dynamicUpdateInterval = 0.5f;
    

    private Native2DArray<float> integrationField;
    private Native2DArray<float2> flowField;
    private Native2DArray<bool> staticObstacleGrid;

    private NativeList<CellPriority> priorityQueueNative;
    
    private int2 destination;
    private bool isInitialized = false;
    private float lastUpdateTime;
    private bool isDirty = false;

    public int2 Destination => destination;
    public float CellSize => cellSize;
    public bool IsInitialized => isInitialized;

    private void Start()
    {
        InitializeGrid();
        SetDestination(new int2(gridWidth / 2, gridHeight / 2));
    }
    
    private void OnDestroy()
    {
        integrationField.Dispose();
        flowField.Dispose();
        staticObstacleGrid.Dispose();
        if (priorityQueueNative.IsCreated)
        {
            priorityQueueNative.Dispose();
        }
    }

    private void Update()
    {
        if (Time.time - lastUpdateTime > dynamicUpdateInterval && isDirty)
        {
            RecalculateField();
            lastUpdateTime = Time.time;
            isDirty = false;
        }
    }

    private void InitializeGrid()
    {
        // 创建托管数组
        integrationField = new Native2DArray<float>(gridWidth, gridHeight, Allocator.Persistent);
        flowField = new Native2DArray<float2>(gridWidth, gridHeight, Allocator.Persistent);
        staticObstacleGrid = new Native2DArray<bool>(gridWidth, gridHeight, Allocator.Persistent);
        
        priorityQueueNative = new NativeList<CellPriority>(Allocator.Persistent);

        // 检测障碍物
        for (int x = 0; x < gridWidth; x++)
        {
            for (int y = 0; y < gridHeight; y++)
            {
                Vector3 worldPos = GridToWorldPosition(new int2(x, y));
                staticObstacleGrid[x, y] = Physics.CheckBox(
                    worldPos, 
                    Vector3.one * cellSize * 0.45f, 
                    Quaternion.identity, 
                    obstacleLayer
                );
            }
        }

        isInitialized = true;
    }
    
    public void SetDestination(int2 newDestination)
    {
        destination = newDestination;
        RecalculateField();
    }

    public void MarkCellsDirty(int2 center, int radius)
    {
        if (!enableDynamicObstacles)
        {
            return;
        }
        
        radius = Mathf.Max(1, radius);
        
        for (int x = -radius; x <= radius; x++)
        {
            for (int y = -radius; y <= radius; y++)
            {
                var gridPos = new int2(center.x + x, center.y + y);
                if (IsPositionValid(gridPos))
                {
                    isDirty = true;
                    return;
                }
            }
        }
    }

    public void RecalculateField()
    {
        if (useJobs)
        {
            CalculateWithJobs();
        }
        else
        {
            CalculateWithoutJobs();
        }
    }

    private void CalculateWithJobs()
    {
        if (!isInitialized)
        {
            return;
        }

        // 重置优先级队列
        priorityQueueNative.Clear();
        
        var dynamicObstaclesNative = DynamicObstacleManager.Instance.GetDynamicObstaclesNative();
        
        // 创建集成场计算作业
        var integrationJob = new FlowFieldIntegrationJob
        {
            GridWidth = gridWidth,
            GridHeight = gridHeight,
            IntegrationField = integrationField,
            StaticObstacleGrid = staticObstacleGrid,
            DynamicObstacleGrid = dynamicObstaclesNative,
            Destination = destination,
            PriorityQueue = priorityQueueNative
        }.Schedule();

        // 创建流场计算作业
        var flowFieldJob = new FlowFieldCalculationJob
        {
            IntegrationField = integrationField,
            StaticObstacleGrid = staticObstacleGrid,
            DynamicObstacleGrid = dynamicObstaclesNative,
            FlowField = flowField,
            GridWidth = gridWidth,
            GridHeight = gridHeight,
            Destination = destination
        };
        
        // 调度并行作业
        var flowFieldHandle = flowFieldJob.Schedule(gridWidth * gridHeight, jobBatchSize, integrationJob);
        flowFieldHandle.Complete();
    }

    private void CalculateWithoutJobs()
    {
        // 传统计算方法（不使用Job System）
        CalculateIntegrationField();
        CalculateFlowField();
    }

    private void CalculateIntegrationField()
    {
        // 初始化集成场
        for (int x = 0; x < gridWidth; x++)
        {
            for (int y = 0; y < gridHeight; y++)
            {
                integrationField[x, y] = float.MaxValue;
            }
        }

        // 目标点成本为0
        integrationField[destination.x, destination.y] = 0;

        // 使用列表作为优先级队列
        var cellsToProcess = new List<int2>();
        cellsToProcess.Add(destination);

        while (cellsToProcess.Count > 0)
        {
            // 找到最小成本的单元格
            int minIndex = 0;
            float minCost = integrationField[cellsToProcess[0].x, cellsToProcess[0].y];
            for (int i = 1; i < cellsToProcess.Count; i++)
            {
                var cell1 = cellsToProcess[i];
                float cost = integrationField[cell1.x, cell1.y];
                if (cost < minCost)
                {
                    minCost = cost;
                    minIndex = i;
                }
            }
            
            var cell = cellsToProcess[minIndex];
            cellsToProcess.RemoveAt(minIndex);
            float currentCost = integrationField[cell.x, cell.y];

            // 检查所有邻居
            for (int dx = -1; dx <= 1; dx++)
            {
                for (int dy = -1; dy <= 1; dy++)
                {
                    if (dx == 0 && dy == 0) continue;

                    var neighbor = new int2(cell.x + dx, cell.y + dy);

                    // 检查邻居是否在网格范围内
                    if (!IsPositionValid(neighbor)) continue;
                    
                    // 跳过障碍物（静态或动态）
                    if (IsAnyObstacle(neighbor)) continue;

                    // 计算移动成本（对角线成本更高）
                    float moveCost = (dx != 0 && dy != 0) ? 1.414f : 1f;
                    float newCost = currentCost + moveCost;

                    // 如果找到更低的成本路径
                    if (newCost < integrationField[neighbor.x, neighbor.y])
                    {
                        integrationField[neighbor.x, neighbor.y] = newCost;
                        cellsToProcess.Add(neighbor);
                    }
                }
            }
        }
    }

    private void CalculateFlowField()
    {
        for (int x = 0; x < gridWidth; x++)
        {
            for (int y = 0; y < gridHeight; y++)
            {
                var gridPos = new int2(x, y);
                
                // 跳过障碍物和目标点
                if (IsAnyObstacle(gridPos) || (x == destination.x && y == destination.y))
                {
                    flowField[x, y] = float2.zero;
                    continue;
                }

                float minCost = float.MaxValue;
                Vector2 bestDirection = Vector2.zero;

                // 检查所有邻居
                for (int dx = -1; dx <= 1; dx++)
                {
                    for (int dy = -1; dy <= 1; dy++)
                    {
                        if (dx == 0 && dy == 0) continue;
                        
                        int nx = x + dx;
                        int ny = y + dy;
                        var neighbor = new int2(nx, ny);

                        // 跳过无效邻居
                        if (!IsPositionValid(neighbor)) continue;
                        if (IsAnyObstacle(neighbor)) continue;

                        // 获取邻居成本
                        float neighborCost = integrationField[nx, ny];
                        
                        // 如果找到更好的方向
                        if (neighborCost < minCost)
                        {
                            minCost = neighborCost;
                            bestDirection = new Vector2(dx, dy).normalized;
                        }
                    }
                }

                flowField[x, y] = bestDirection;
            }
        }
    }

    public float2 GetFlowDirection(Vector3 worldPosition)
    {
        var gridPos = WorldToGridPosition(worldPosition);
        if (!IsPositionValid(gridPos)) return Vector2.zero;
        
        return flowField[gridPos.x, gridPos.y];
    }

    public int2 WorldToGridPosition(Vector3 worldPosition)
    {
        int x = Mathf.FloorToInt((worldPosition.x - transform.position.x) / cellSize + gridWidth / 2);
        int y = Mathf.FloorToInt((worldPosition.z - transform.position.z) / cellSize + gridHeight / 2);
        return new int2(
            Mathf.Clamp(x, 0, gridWidth - 1),
            Mathf.Clamp(y, 0, gridHeight - 1)
        );
    }

    public Vector3 GridToWorldPosition(int2 gridPosition)
    {
        float x = (gridPosition.x - gridWidth / 2) * cellSize + cellSize / 2;
        float z = (gridPosition.y - gridHeight / 2) * cellSize + cellSize / 2;
        return new Vector3(x, 0, z) + transform.position;
    }

    public bool IsPositionValid(int2 gridPos)
    {
        return gridPos.x >= 0 && gridPos.x < gridWidth && 
               gridPos.y >= 0 && gridPos.y < gridHeight;
    }

    public bool IsStaticObstacle(int2 gridPos)
    {
        return staticObstacleGrid[gridPos.x, gridPos.y];
    }
    
    public bool IsDynamicObstacle(int2 gridPos)
    {
        if (DynamicObstacleManager.Instance != null)
        {
            return DynamicObstacleManager.Instance.IsDynamicObstacle(gridPos);
        }
        return false;
    }
    
    public bool IsAnyObstacle(int2 gridPos)
    {
        return IsStaticObstacle(gridPos) || IsDynamicObstacle(gridPos);
    }

    private void OnDrawGizmos()
    {
        if (!Application.isPlaying || !isInitialized) return;

        // 绘制网格单元格
        for (int x = 0; x < gridWidth; x++)
        {
            for (int y = 0; y < gridHeight; y++)
            {
                Vector3 center = GridToWorldPosition(new int2(x, y));
                
                // 绘制静态障碍物
                if (staticObstacleGrid[x, y])
                {
                    Gizmos.color = Color.red;
                    Gizmos.DrawCube(center, Vector3.one * cellSize * 0.9f);
                }
                // 绘制动态障碍物
                else if (DynamicObstacleManager.Instance != null && DynamicObstacleManager.Instance.IsDynamicObstacle(new int2(x, y)))
                {
                    Gizmos.color = new Color(1, 0.5f, 0, 0.5f); // 橙色
                    Gizmos.DrawCube(center, Vector3.one * cellSize * 0.9f);
                }
                // 绘制目标点
                else if (x == destination.x && y == destination.y)
                {
                    Gizmos.color = Color.green;
                    Gizmos.DrawCube(center, Vector3.one * cellSize * 0.9f);
                }
                // 绘制普通网格单元格
                else
                {
                    Gizmos.color = Color.gray;
                    Gizmos.DrawWireCube(center, Vector3.one * cellSize);
                }
            }
        }

        // 绘制流场方向
        for (int x = 0; x < gridWidth; x++)
        {
            for (int y = 0; y < gridHeight; y++)
            {
                if (IsAnyObstacle(new int2(x, y)) || math.all(flowField[x, y] == float2.zero))
                {
                    continue;
                }

                Vector3 start = GridToWorldPosition(new int2(x, y));
                Vector3 end = start + new Vector3(flowField[x, y].x, 0, flowField[x, y].y) * cellSize * 0.5f;
                
                // 根据成本设置颜色
                float normalizedCost = Mathf.Clamp01(integrationField[x, y] / 100f);
                Color color = Color.Lerp(Color.green, Color.red, normalizedCost);
                
                Gizmos.color = color;
                Gizmos.DrawLine(start, end);
                
                // 绘制箭头头部
                GizmoUtils.DrawArrow(start, end, color, cellSize * 0.2f);
            }
        }
    }
}
